
module(..., package.seeall) 

local function parent (i) 
   if i>1
   then return math.floor(i/2) 
   else return nil
   end
end
   
local function left_child (i)
   return i*2
end 

local function right_child (i)
   return (i*2)+1 
end 

local function make_empty_heap(above, locate) 
   return { pref = above, 
	    loc = locate }
end

--[[

heapify_up(h,i) takes an almost-heap 
h, that is, an array in which for any 
other index j, h[j] belongs below its 
parent. 

If h[i] belongs *above* its parent, 
then heapify_up does this interchange, 
and then continues to
heapify_up(h,parent(i)).  

--]]

function heapify_up(h,i)
   local p = parent(i) 
   if p and h[i] and 
      (not(h[p]) or h.pref(h[i],h[p]))
   then 
      -- flip em!
      h[p], h[i] = h[i], h[p]
      h.loc(h[p],p)
      h.loc(h[i],i)
      heapify_up(h,p)
   end 
   check_heap(h)
end


--[[

heapify_down takes a broken heap and 
takes one step to repair it.  In particular, 
its arguments are an array a and an index 
i, such that: 

1.  the left and right subtrees dominated by i 
    are already heaps; 

2.  a[i] belongs below its parent if any.  

(Note:  "left and right subtrees below i" 
	is a little tricky to define precisely!)  

If a[i] is less than either of its children, 
heapify_down flips it with the larger of the 
two children.  It recursively repairs any 
damage this may have caused to the heap property 
of the child.  

--]]

local function heapify_down (h,i) 
   local l = left_child(i)
   local r = right_child(i)
   local f = nil 

   if h[l] and h[r]
   then
      if h.pref(h[l],h[r]) then f=l else f=r end 
   else 
      if h[l] then f=l else 
	 if h[r] then f=r end end
   end
   
   if f then
      if not(h[i]) or h.pref(h[f],h[i]) 
      -- if the favored child 
      -- belongs on top, then... 
      then 
	 -- flip em! 
	 h[f], h[i] = h[i], h[f]
	 -- and record location updates 
	 h.loc(h[f],f)
	 h.loc(h[i],i)
	 -- recurse down the affected path 
	 heapify_down (h,f)
      end
   end 
   check_heap(h)   
end  




--[[

build_heap takes any array of numbers and 
optionally functions for comparison and 
recording updates.  

build_heap initializes a heap in which each
 parent > its children.
Thus, the root is the max element.  

--]]

function build_heap (a, f, g) 
   local above = f or
      function (a,b) return a>b end 

   local locate = g or
      function (a,b) end

   local h =  make_empty_heap(above, locate) 

   for i = 1,#a do 
      h[i]=a[i]; locate(a[i],i)
   end 

   for i = math.ceil((#h)/2), 1, -1 do 
      heapify_down(h,i) 
   end 
   return h 
end 

function build_heap_up (a, f, g) 
   local above = f or 
      function (a,b) return a>b end 

   local locate = g or 
      function (a,b) end

   local h =  make_empty_heap(above, locate) 

   for i = 1,#a do 
      h[i]=a[i] 
      h.loc(a[i],i)
      heapify_up(h,i)
   end 
   return h 
end 

function insert(h,v) 
   local i = #h+1 
   h[i] = v 			-- place v in new spot
   h.loc(v,i)			-- inform the client 
   return heapify_up(h,i)
end 

function find_top(h) 
   return h[1]
end 

function delete(h,i) 
   h.loc(h[i],nil)		-- the old h[i] will be nowhere 
   h.loc(nil,i)			-- nil will be in spot i
   h[i] = nil 
   heapify_down(h,i)		-- fix h if broken at i 
   check_heap(h)
   return h_new 
end 

function extract_top(h) 
   local v = h[1]		-- find the top value 
   delete(h,1) 			-- delete it
   return v 			-- return the old top dog 
end 

function check_heap (h) 
   for loc,val in pairs(h) do 
      if val and type(loc) == number 
      then 
	 local l = left_child(loc) 
	 local r = right_child(loc)
	 if (h[l]) and h.pref(h[l], h[loc])
	 then 
	    error(string.format("check_heap failed: %d vs %d\n", loc,l))
	 else if 
	    (h[r]) and h.pref(h[r],h[loc])
	    then 			
	       error(string.format("check_heap failed: %d vs %d\n", loc,r))
	    end 
	 end  
      end 
   end 
end 
